# -*-coding:utf8-*-
from tools.search_paper import get_paper_df
from tools.MultiPaperAnalysis import MultiPaper
from Searcher.llm import llm
from tools.template_utils import Template
import pandas as pd


def write_paper(title, lang='en',state={},is_title=True):
    if not is_title:
        prompt = f'''
You will call a summary writing tool with the input being an English summary title. Based on the user's needs, only output an English summary title. If the user's input is already an English summary title, do not make any changes and output it as is.
User's need: {title}
Please only output the English title:
'''
        title = llm(prompt)
    # 搜索并总结，筛选文档
    paper_info_df = get_paper_df(title)
    paper_info_df['ID'] = range(len(paper_info_df))
    paper_info_df['ID'] += 1
    print('文档结束')
    #  开始写作
    mp = MultiPaper(title, paper_info_df, llm,lang=lang)
    #  随机选择写作模板
    t = Template(lang)

    import concurrent.futures

    # 定义两个函数的包装器，以便在线程池中执行
    def write_introduction():
        return mp.write_introduction(t.random_template_yy())

    def write_methods():
        return mp.write_methods(t.random_template_ff())

    # 创建一个包含两个线程的线程池
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        # 提交两个任务
        future_introduction = executor.submit(write_introduction)
        future_methods = executor.submit(write_methods)

        # 等待两个任务完成并获取结果
        state['引言'] = future_introduction.result()
        state['方法'] = future_methods.result()
    # introduction 和 methods 现在都已经计算完成

    state['讨论'] = mp.write_discussion(t.random_template_tl())
    state['结论'] = mp.write_conclusion(t.random_template_jl())

    # 增加引用对照
    references = paper_info_df['Title'].tolist()

    references_md = ''

    return_md = f"# Introduction\n\n{state['引言']}\n\n# Methods\n\n{state['方法']}\n" + \
                f"\n# Discussion\n\n{state['讨论']}\n\n# Conclusion\n\n{state['结论']}\n\n"


    # 判断那些文章被引用了
    total = 0
    references_list= []
    for i, j in enumerate(references):
        if f"<sup>{i + 1}</sup>" in return_md:
            references_md += f'[{i + 1}] {j}\n\n'
            references_list.append(i+1)
            total += 1
    # 给文章重新排序
    for i in range(len(references_list)):
        r_id = references_list[i]
        references_md = references_md.replace(f"[{r_id}]",f"[{i+1}]")
        return_md = return_md.replace(f"<sup>{r_id}</sup>",f"<sup>{i+1}</sup>")
    return_md += f"# References\n\n{references_md}\n"

    return return_md

if __name__ == '__main__':
    import time
    start = time.time()
    md_content = write_paper('基于知识图谱的检索增强生成（GraphRAG）',lang='zh',is_title=False)
    print(time.time()-start)
    with open('result.md', 'w', encoding='utf-8') as f:
        f.write(md_content)

    import hashlib

    '6e9e992fa80b8d9bcb54722acfae2579dcf76c64'
    def calculate_file_hash(file_path):
        hash_func = getattr(hashlib, "md5")()

        with open(file_path, "rb") as f:
            while chunk := f.read(8192):
                hash_func.update(chunk)

        return hash_func.hexdigest()

    print(calculate_file_hash('result.md'))

    # import argparse
    #
    # parser = argparse.ArgumentParser(description='Write a paper with a given title.')
    # parser.add_argument('--topic', type=str, required=True, help='The title of the paper')
    # args = parser.parse_args()
    # title = args.topic
    # write_paper(title)
